#pragma once
#include <vector>
#include <memory>
#include <algorithm>
#include <Eigen/Dense>

/**
 * @brief KD树节点
 */
struct KDNode {
    Eigen::VectorXd point;     // 节点坐标
    size_t index;              // 原始点的索引
    int axis;                  // 分割轴
    std::shared_ptr<KDNode> left;   // 左子树
    std::shared_ptr<KDNode> right;  // 右子树
    
    KDNode(const Eigen::VectorXd& p, size_t idx, int ax) 
        : point(p), index(idx), axis(ax) {}
};

/**
 * @brief KD树实现
 */
class KDTree {
public:
    /**
     * @brief 构建KD树
     * @param points 输入点集
     */
    explicit KDTree(const Eigen::MatrixXd& points);
    
    /**
     * @brief 查找最近邻
     * @param query 查询点
     * @return 最近点的索引和距离
     */
    std::pair<size_t, double> nearest(const Eigen::VectorXd& query) const;
    
    /**
     * @brief 查找K个最近邻
     * @param query 查询点
     * @param k 邻居数量
     * @param indices 返回的邻居索引
     * @param distances 返回的距离
     */
    void kNearest(const Eigen::VectorXd& query, int k,
                 std::vector<size_t>& indices,
                 std::vector<double>& distances) const;

private:
    std::shared_ptr<KDNode> buildTree(const std::vector<size_t>& indices,
                                    int depth);
                                    
    void nearestImpl(const Eigen::VectorXd& query,
                    const std::shared_ptr<KDNode>& node,
                    size_t& bestIndex,
                    double& bestDist) const;
                    
    void kNearestImpl(const Eigen::VectorXd& query,
                     const std::shared_ptr<KDNode>& node,
                     std::vector<std::pair<double, size_t>>& nearest,
                     int k) const;

    Eigen::MatrixXd points_;  // 所有点的坐标
    std::shared_ptr<KDNode> root_;  // 树的根节点
    int dim_;  // 空间维度
}; 